# in app/schemas/response.py
from datetime import datetime
from typing import TypeVar, Generic, Optional, Any, Dict
import pytz
from pydantic import BaseModel, Field, model_serializer

T = TypeVar('T')


class BaseSchema(BaseModel):
    """Base schema with timezone conversion for datetime fields"""

    @model_serializer
    def serialize_model(self) -> Dict[str, Any]:
        """Convert all datetime fields from UTC to Asia/Shanghai and format as 'YYYY-MM-DD HH:MM:SS'"""
        data = {}

        # Get China timezone
        china_tz = pytz.timezone('Asia/Shanghai')

        for field_name, field_value in self.__dict__.items():
            if isinstance(field_value, datetime):
                # Convert UTC to China timezone and format
                if field_value.tzinfo is None:
                    # Assume UTC if no timezone info
                    utc_time = pytz.UTC.localize(field_value)
                else:
                    utc_time = field_value.astimezone(pytz.UTC)

                china_time = utc_time.astimezone(china_tz)
                data[field_name] = china_time.strftime('%Y-%m-%d %H:%M:%S')
            else:
                data[field_name] = field_value

        return data


class UnifiedResponse(BaseModel, Generic[T]):
    code: int = Field(default=200, description="状态码")
    message: str = Field(default="操作成功", description="响应消息")
    data: Optional[T] = None


# 可以预定义一些常用的成功响应
class SuccessResponse(UnifiedResponse):
    pass


# 或者一个创建成功响应的函数
def success_response(data: T, msg: str = "操作成功") -> UnifiedResponse[T]:
    return UnifiedResponse(data=data, message=msg)
